#!/bin/python3
import struct

def trans2jxpamg_matrix(matrix_dir, save_matrix):
    f = open(matrix_dir, "rb")

    numrow = f.read(4)
    numrow = struct.unpack("i", numrow)[0]
    nnz    = f.read(4)
    nnz = struct.unpack("i", nnz)[0]
    print(nnz)

    for i in range(0, numrow):
        data = f.read(4)

    row_ptr = []
    for i in range(0, numrow + 1 ):
        data = f.read(4)
        data = struct.unpack("i", data)[0]
        row_ptr.append(data)

    if(nnz != row_ptr[-1]):
        print("nnz is {0}, row_ptr[-1] is {1}".format(str(nnz), str(row_ptr[-1])))
        exit(0)

    col = []
    for i in range(0, row_ptr[-1]):
        data = f.read(4)
        data = struct.unpack("i", data)[0]
        col.append(data)

    val = []
    for i in range(0, row_ptr[-1]):
        data = f.read(8)
        data = str(struct.unpack("d", data)[0])
        val.append(data)

    f.close()
    print("\nread matrix success!\n")

    with open(save_matrix, "w+" ) as fo:
        fo.write(str(numrow)+"\n")
        for i in range(0, numrow + 1):
            fo.write(str(row_ptr[i])+"\n")
        for i in range(0, len(col)):
            fo.write(str(col[i])+"\n")
        for i in range(0, len(val)):
            fo.write(str(val[i])+"\n")

def trans2jxpamg_rhs(rhs_dir, save_rhs):
    f = open(rhs_dir, "rb")
    numrow = f.read(4)
    numrow = struct.unpack("i", numrow)[0]

    rhs = []
    for i in range(numrow):
        data = f.read(4)
        data = f.read(8)
        data = str(struct.unpack("d", data)[0])
        rhs.append(data)

    f.close()

    print("\n read rhs success!\n")

    with open(save_rhs, "w+" ) as fo:
        fo.write(str(numrow)+"\n")
        for i in range(0, numrow):
            fo.write(str(rhs[i])+"\n")


if __name__=="__main__":
    matrix_dir  = "/home/matrix/hsl_matrix/mat.dat"
    rhs_dir     = "/home/matrix/hsl_matrix/rhs.dat"
    save_matrix = '/home/matrix/hsl_matrix/jxpamg_matrix'
    save_rhs = "/home/matrix/hsl_matrix/jxpamg_rhs"
    trans2jxpamg_matrix(matrix_dir, save_matrix)
    trans2jxpamg_rhs(rhs_dir, save_rhs)

